import numpy as np
import pandas as pd


def from_open3d(o3d_data, **kwargs):
    """Create a PyntCloud instance from Open3D's PointCloud/TriangleMesh instance"""
    try:
        import open3d as o3d
    except ImportError:
        raise ImportError("Open3D must be installed. Try `pip install open3d`")

    if not isinstance(o3d_data, (o3d.geometry.PointCloud, o3d.geometry.TriangleMesh)):
        raise TypeError(
            f"Type {type(o3d_data)} not supported for conversion."
            f"Expected {o3d.geometry.PointCloud} or {o3d.geometry.TriangleMesh}"
        )

    mesh = None
    if isinstance(o3d_data, o3d.geometry.TriangleMesh):
        mesh = pd.DataFrame(
            data=np.asarray(o3d_data.triangles), columns=["v1", "v2", "v3"]
        )

        points = pd.DataFrame(
            data=np.asarray(o3d_data.vertices), columns=["x", "y", "z"]
        )

        if o3d_data.vertex_colors:
            colors = (np.asarray(o3d_data.vertex_colors) * 255).astype(np.uint8)
            points["red"] = colors[:, 0]
            points["green"] = colors[:, 1]
            points["blue"] = colors[:, 2]

        if o3d_data.vertex_normals:
            normals = np.asarray(o3d_data.vertex_normals)
            points["nx"] = normals[:, 0]
            points["ny"] = normals[:, 1]
            points["nz"] = normals[:, 2]

    elif isinstance(o3d_data, o3d.geometry.PointCloud):
        points = pd.DataFrame(data=np.asarray(o3d_data.points), columns=["x", "y", "z"])

        if o3d_data.colors:
            colors = (np.asarray(o3d_data.colors) * 255).astype(np.uint8)
            points["red"] = colors[:, 0]
            points["green"] = colors[:, 1]
            points["blue"] = colors[:, 2]

        if o3d_data.normals:
            normals = np.asarray(o3d_data.normals)
            points["nx"] = normals[:, 0]
            points["ny"] = normals[:, 1]
            points["nz"] = normals[:, 2]

    return {"points": points, "mesh": mesh}


def to_open3d(cloud, mesh=True, colors=True, normals=True, **kwargs):
    """Convert PyntCloud's instance `cloud` to Open3D's PointCloud/TriangleMesh instance"""
    try:
        import open3d as o3d
    except ImportError:
        raise ImportError("Open3D must be installed. Try `pip install open3d`")

    if mesh and cloud.mesh is not None:
        triangle_mesh = o3d.geometry.TriangleMesh()
        triangle_mesh.triangles = o3d.utility.Vector3iVector(
            cloud.mesh[["v1", "v2", "v3"]].values
        )
        triangle_mesh.vertices = o3d.utility.Vector3dVector(cloud.xyz)
        if colors and {"red", "green", "blue"}.issubset(cloud.points.columns):
            triangle_mesh.vertex_colors = o3d.utility.Vector3dVector(
                cloud.points[["red", "green", "blue"]].values
            )
        if normals and {"nx", "ny", "nz"}.issubset(cloud.points.columns):
            triangle_mesh.vertex_normals = o3d.utility.Vector3dVector(
                cloud.points[["nx", "ny", "nz"]].values
            )
        return triangle_mesh
    else:
        point_cloud = o3d.geometry.PointCloud()
        point_cloud.points = o3d.utility.Vector3dVector(cloud.xyz)
        if colors and {"red", "green", "blue"}.issubset(cloud.points.columns):
            point_cloud.colors = o3d.utility.Vector3dVector(
                cloud.points[["red", "green", "blue"]].values
            )
        if normals and {"nx", "ny", "nz"}.issubset(cloud.points.columns):
            point_cloud.normals = o3d.utility.Vector3dVector(
                cloud.points[["nx", "ny", "nz"]].values
            )
        return point_cloud
